# In this problem you will implement SLAM in a 2 dimensional
# world. Please define a function, slam, which takes five
# parameters as input and returns the vector mu. This vector
# should have x, y coordinates interlaced, so for example, 
# if there were 2 poses and 2 landmarks, mu would look like:
#
#  mu =  matrix([[Px0],
#                [Py0],
#                [Px1],
#                [Py1],
#                [Lx0],
#                [Ly0],
#                [Lx1],
#                [Ly1]])
#
# data - This is the data that is generated with the included
#        make_data function. You can also use test_data to
#        make sure your function gives the correct result.
#
# N -    The number of time steps.
#
# num_landmarks - The number of landmarks.
#
# motion_noise - The noise associated with motion. The update
#                strength for motion should be 1.0 / motion_noise.
#
# measurement_noise - The noise associated with measurement.
#                     The update strength for measurement should be
#                     1.0 / measurement_noise.
import random
from math import *

from matrix import matrix
from robot import robot, make_data, print_result, check_mu


# --------------------------------
#
# slam - retains entire path and all landmarks
#
def slam(data, N, num_landmarks, world_size, motion_noise, measurement_noise):
    # Set the dimension of the filter
    dim = 2 * (N + num_landmarks)
    
    # make the constraint information matrix and vector
    Omega = matrix()
    Omega.zero(dim, dim)
    Omega.value[0][0] = 1.0
    Omega.value[1][1] = 1.0
    
    Xi = matrix()
    Xi.zero(dim, 1)
    Xi.value[0][0] = world_size / 2.0
    Xi.value[1][0] = world_size / 2.0
    
    # process the data
    for k in range(len(data)):
        # n is the index of the robot pose in the matrix/vector
        n = k * 2 
        
        measurement = data[k][0]
        motion      = data[k][1]
        
        # integrate the measurements
        for i in range(len(measurement)):
            # m is the index of the landmark coordinate in the matrix/vector
            m = 2 * (N + measurement[i][0])
            
            # update the information maxtrix/vector based on the measurement
            for b in range(2):
                Omega.value[n+b][n+b] +=  1.0 / measurement_noise
                Omega.value[m+b][m+b] +=  1.0 / measurement_noise
                Omega.value[n+b][m+b] += -1.0 / measurement_noise
                Omega.value[m+b][n+b] += -1.0 / measurement_noise
                Xi.value[n+b][0]      += -measurement[i][1+b] / measurement_noise
                Xi.value[m+b][0]      +=  measurement[i][1+b] / measurement_noise
        
        # update the information maxtrix/vector based on the robot motion
        for b in range(4):
            Omega.value[n+b][n+b] +=  1.0 / motion_noise
        for b in range(2):
            Omega.value[n+b  ][n+b+2] += -1.0 / motion_noise
            Omega.value[n+b+2][n+b  ] += -1.0 / motion_noise
            Xi.value[n+b  ][0]        += -motion[b] / motion_noise
            Xi.value[n+b+2][0]        +=  motion[b] / motion_noise
    
    # compute best estimate
    mu = Omega.inverse() * Xi
    return mu

#
# Main routine
#
if __name__ == '__main__':
    num_landmarks      = 5        # number of landmarks
    N                  = 20       # time steps
    world_size         = 100.0    # size of world
    measurement_range  = 50.0     # range at which we can sense landmarks
    motion_noise       = 2.0      # noise in robot motion
    measurement_noise  = 2.0      # noise in the measurements
    distance           = 20.0     # distance by which robot (intends to) move each iteratation 
    
    data = make_data(N, num_landmarks, world_size, measurement_range, motion_noise, measurement_noise, distance)
    result = slam(data, N, num_landmarks, world_size, motion_noise, measurement_noise)
    
    # -------------
    # Testing
    
    ##  Test Case 1
    answer_mu = matrix([
        ##  Estimated Pose(s):
        [49.998897453299165], [49.998505706587814],
        [37.9714477943764  ], [33.64993445823509 ],
        [26.183449863995392], [18.153338459791925],
        [13.743443839776688], [2.1141193319706257],
        [28.095060682659934], [16.78089653056425 ],
        [42.382654337758865], [30.899617637854934],
        [55.82918373374959 ], [44.494287384838586],
        [70.85548928962663 ], [59.69712516287841 ],
        [85.6953531635832  ], [75.54007207500423 ],
        [74.00974406829611 ], [92.43147558585063 ],
        [53.54264788322474 ], [96.45111370814985 ],
        [34.52341231228876 ], [100.07762713840204],
        [48.621309970082486], [83.95097871134821 ],
        [60.19521941022714 ], [68.1050904555393  ],
        [73.77592978594885 ], [52.932451315943574],
        [87.12997140410576 ], [38.53576961176431 ],
        [80.3007799395094  ], [20.505859780712917],
        [72.79656231764604 ], [2.942675607797428 ],
        [55.243590482706026], [13.253021809459227],
        [37.41439688492312 ], [22.31502245240636 ],
        
        ##  Estimated Landmarks:
        [82.95425645256424 ], [13.536536707427121],
        [70.49306062345174 ], [74.13913606764582 ],
        [36.73812342335858 ], [61.2789905549488  ],
        [18.696326102039485], [66.05733561281015 ],
        [20.632945056999347], [16.87255837889543]
    ])
    test_data1 = [[[[1, 19.457599255548065, 23.8387362100849], [2, -13.195807561967236, 11.708840328458608], [3, -30.0954905279171, 15.387879242505843]], [-12.2607279422326, -15.801093326936487]], [[[2, -0.4659930049620491, 28.088559771215664], [4, -17.866382374890936, -16.384904503932]], [-12.2607279422326, -15.801093326936487]], [[[4, -6.202512900833806, -1.823403210274639]], [-12.2607279422326, -15.801093326936487]], [[[4, 7.412136480918645, 15.388585962142429]], [14.008259661173426, 14.274756084260822]], [[[4, -7.526138813444998, -0.4563942429717849]], [14.008259661173426, 14.274756084260822]], [[[2, -6.299793150150058, 29.047830407717623], [4, -21.93551130411791, -13.21956810989039]], [14.008259661173426, 14.274756084260822]], [[[1, 15.796300959032276, 30.65769689694247], [2, -18.64370821983482, 17.380022987031367]], [14.008259661173426, 14.274756084260822]], [[[1, 0.40311325410337906, 14.169429532679855], [2, -35.069349468466235, 2.4945558982439957]], [14.008259661173426, 14.274756084260822]], [[[1, -16.71340983241936, -2.777000269543834]], [-11.006096015782283, 16.699276945166858]], [[[1, -3.611096830835776, -17.954019226763958]], [-19.693482634035977, 3.488085684573048]], [[[1, 18.398273354362416, -22.705102332550947]], [-19.693482634035977, 3.488085684573048]], [[[2, 2.789312482883833, -39.73720193121324]], [12.849049222879723, -15.326510824972983]], [[[1, 21.26897046581808, -10.121029799040915], [2, -11.917698965880655, -23.17711662602097], [3, -31.81167947898398, -16.7985673023331]], [12.849049222879723, -15.326510824972983]], [[[1, 10.48157743234859, 5.692957082575485], [2, -22.31488473554935, -5.389184118551409], [3, -40.81803984305378, -2.4703329790238118]], [12.849049222879723, -15.326510824972983]], [[[0, 10.591050242096598, -39.2051798967113], [1, -3.5675572049297553, 22.849456408289125], [2, -38.39251065320351, 7.288990306029511]], [12.849049222879723, -15.326510824972983]], [[[0, -3.6225556479370766, -25.58006865235512]], [-7.8874682868419965, -18.379005523261092]], [[[0, 1.9784503557879374, -6.5025974151499]], [-7.8874682868419965, -18.379005523261092]], [[[0, 10.050665232782423, 11.026385307998742]], [-17.82919359778298, 9.062000642947142]], [[[0, 26.526838150174818, -0.22563393232425621], [4, -33.70303936886652, 2.880339841013677]], [-17.82919359778298, 9.062000642947142]]]
    result = slam(test_data1, 20, 5, world_size, 2.0, 2.0)
    print_result(20, 5, result)
    check_mu(result, answer_mu)
    
    ##  Test Case 2
    answer_mu = matrix([
        ##  Estimated Pose(s):
        [49.999477332348086], [49.99890156551778 ],
        [69.18030735650325 ], [45.66363344648087 ],
        [87.7418064828167  ], [39.7015808043011  ],
        [76.2689002984748  ], [56.308708837988604],
        [64.31595832940498 ], [72.17416433180664 ],
        [52.25593780877982 ], [88.15129488583568 ],
        [44.05788183012242 ], [69.39888197470297 ],
        [37.00060088117602 ], [49.91574097036883 ],
        [30.923168865968172], [30.953132613521404],
        [23.507016577616362], [11.417418859810084],
        [34.17904823546437 ], [27.13111103387763 ],
        [44.153991513067325], [43.8439571694372  ],
        [54.80488988167458 ], [60.918509500376814],
        [65.69693171880341 ], [78.54360707650197 ],
        [77.4671073715615  ], [95.62410423059377 ],
        [96.80062212965419 ], [98.8189438054247  ],
        [75.95566185322726 ], [99.96919845616313 ],
        [70.19856001872532 ], [81.17887019996131 ],
        [64.05274158851972 ], [61.72101680774857 ],
        [58.10609891363755 ], [42.62553459605071 ],
        
        ##  Estimated Landmarks:
        [76.7777215386166  ], [42.88538378284734 ],
        [85.06362510793153 ], [77.43622580258783 ],
        [13.546262775052924], [95.64892318002514 ],
        [59.447682292240444], [39.59345703309377 ],
        [69.26225738002688 ], [94.23786125406279 ]
    ])
    test_data2 = [[[[0, 26.543274387283322, -6.262538160312672], [3, 9.937396825799755, -9.128540360867689]], [18.92765331253674, -6.460955043986683]], [[[0, 7.706544739722961, -3.758467215445748], [1, 17.03954411948937, 31.705489938553438], [3, -11.61731288777497, -6.64964096716416]], [18.92765331253674, -6.460955043986683]], [[[0, -12.35130507136378, 2.585119104239249], [1, -2.563534536165313, 38.22159657838369], [3, -26.961236804740935, -0.4802312626141525]], [-11.167066095509824, 16.592065417497455]], [[[0, 1.4138633151721272, -13.912454837810632], [1, 8.087721200818589, 20.51845934354381], [3, -17.091723454402302, -16.521500551709707], [4, -7.414211721400232, 38.09191602674439]], [-11.167066095509824, 16.592065417497455]], [[[0, 12.886743222179561, -28.703968411636318], [1, 21.660953298391387, 3.4912891084614914], [3, -6.401401414569506, -32.321583037341625], [4, 5.034079343639034, 23.102207946092893]], [-11.167066095509824, 16.592065417497455]], [[[1, 31.126317672358578, -10.036784369535214], [2, -38.70878528420893, 7.4987265861424595], [4, 17.977218575473767, 6.150889254289742]], [-6.595520680493778, -18.88118393939265]], [[[1, 41.82460922922086, 7.847527392202475], [3, 15.711709540417502, -30.34633659912818]], [-6.595520680493778, -18.88118393939265]], [[[0, 40.18454208294434, -6.710999804403755], [3, 23.019508919299156, -10.12110867290604]], [-6.595520680493778, -18.88118393939265]], [[[3, 27.18579315312821, 8.067219022708391]], [-6.595520680493778, -18.88118393939265]], [[], [11.492663265706092, 16.36822198838621]], [[[3, 24.57154567653098, 13.461499960708197]], [11.492663265706092, 16.36822198838621]], [[[0, 31.61945290413707, 0.4272295085799329], [3, 16.97392299158991, -5.274596836133088]], [11.492663265706092, 16.36822198838621]], [[[0, 22.407381798735177, -18.03500068379259], [1, 29.642444125196995, 17.3794951934614], [3, 4.7969752441371645, -21.07505361639969], [4, 14.726069092569372, 32.75999422300078]], [11.492663265706092, 16.36822198838621]], [[[0, 10.705527984670137, -34.589764174299596], [1, 18.58772336795603, -0.20109708164787765], [3, -4.839806195049413, -39.92208742305105], [4, 4.18824810165454, 14.146847823548889]], [11.492663265706092, 16.36822198838621]], [[[1, 5.878492140223764, -19.955352450942357], [4, -7.059505455306587, -0.9740849280550585]], [19.628527845173146, 3.83678180657467]], [[[1, -11.150789592446378, -22.736641053247872], [4, -28.832815721158255, -3.9462962046291388]], [-19.841703647091965, 2.5113335861604362]], [[[1, 8.64427397916182, -20.286336970889053], [4, -5.036917727942285, -6.311739993868336]], [-5.946642674882207, -19.09548221169787]], [[[0, 7.151866679283043, -39.56103232616369], [1, 16.01535401373368, -3.780995345194027], [4, -3.04801331832137, 13.697362774960865]], [-5.946642674882207, -19.09548221169787]], [[[0, 12.872879480504395, -19.707592098123207], [1, 22.236710716903136, 16.331770792606406], [3, -4.841206109583004, -21.24604435851242], [4, 4.27111163223552, 32.25309748614184]], [-5.946642674882207, -19.09548221169787]]]
    result = slam(test_data2, 20, 5, world_size, 2.0, 2.0)
    print_result(20, 5, result)
    check_mu(result, answer_mu)
